#ifndef CNN_H
#define CNN_H

#include <vector>

#include "ui.h"

extern float X_scaler_mean[8];
extern float X_scaler_scale[8];
extern float Voltage_scaled[8];
extern float y_scaler_mean;
extern float y_scaler_scale;

extern std::vector<std::vector<float>> input;

class MultiTaskModel {
public:
    // conv1_parameters
    std::vector<std::vector<std::vector<float>>> conv1_weights;
    std::vector<float> conv1_bias;

    // shared_parameters
    std::vector<std::vector<float>> shared_weights;
    std::vector<float> shared_bias;

    // regression_head_parameters
    std::vector<std::vector<float>> regression_head_weights;
    std::vector<float> regression_head_bias;

    // classification_head_parameters
    std::vector<std::vector<float>> classification_head_weights;
    std::vector<float> classification_head_bias;

    // forward
    std::pair<float, std::vector<float>> forward(const std::vector<std::vector<float>>& input);

    MultiTaskModel();

    void initWeights();

private:

    float relu1d(float x);

    std::vector<std::vector<float>> relu2d(const std::vector<std::vector<float>>& input);

    std::vector<float> softmax(const std::vector<float>& logits);

    std::vector<float> conv1d_single_channel(const std::vector<std::vector<float>>& input, 
                               const std::vector<std::vector<std::vector<float>>>& weights, 
                               const std::vector<float>& bias, 
                               int stride, int padding, int kernel_size);

    std::vector<std::vector<float>> conv1d(const std::vector<std::vector<float>>& input, 
                                                       const std::vector<std::vector<std::vector<float>>>& weights, 
                                                       const std::vector<float>& bias, 
                                                       int stride, int padding, int kernel_size);

    std::vector<float> max_pool1d(const std::vector<float>& input, int kernel_size);

    std::vector<std::vector<float>> max_pool2d(const std::vector<std::vector<float>>& input, int kernel_size);

    std::vector<float> flatten(const std::vector<std::vector<float>>& input);

    std::vector<float> fully_connected(const std::vector<float>& input, 
                                       const std::vector<std::vector<float>>& weights, 
                                       const std::vector<float>& bias);


};

extern MultiTaskModel model;
void init_variables();
void sliding_window(void);
void model_forward(void);

#endif // CNN_H
